{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xeOG96gXPeqz"
      },
      "source": [
        "# Security Sentinel\n",
        "\n",
        "### Audit and harden your code in one sweep\n",
        "\n",
        "Choose a contract to inspect a snippet for security issues:\n",
        "\n",
        "- Scan for vulnerabilities\n",
        "- Draft threat intelligence notes\n",
        "- Suggest secure fixes\n",
        "- Craft exploit-focused tests"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B7ftYo53Pw94",
        "outputId": "9daa3972-d5a1-4cd2-9952-cd89a54c6ddd"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import logging\n",
        "from enum import StrEnum\n",
        "from getpass import getpass\n",
        "\n",
        "import gradio as gr\n",
        "from openai import OpenAI\n",
        "from dotenv import load_dotenv\n",
        "\n",
        "load_dotenv(override=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AXmPDuydPuUp"
      },
      "outputs": [],
      "source": [
        "logging.basicConfig(level=logging.WARNING)\n",
        "\n",
        "logger = logging.getLogger('sentinel')\n",
        "logger.setLevel(logging.DEBUG)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def get_secret_in_google_colab(env_name: str) -> str:\n",
        "    try:\n",
        "      from google.colab import userdata\n",
        "      return userdata.get(env_name)\n",
        "    except Exception:\n",
        "      return ''\n",
        "\n",
        "\n",
        "def get_secret(env_name: str) -> str:\n",
        "    key = os.environ.get(env_name) or get_secret_in_google_colab(env_name)\n",
        "\n",
        "    if not key:\n",
        "        key = getpass(f'Enter {env_name}:').strip()\n",
        "\n",
        "    if key:\n",
        "        logger.info(f'✅ {env_name} provided')\n",
        "    else:\n",
        "        logger.warning(f'❌ {env_name} not provided')\n",
        "    return key.strip()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "d7Qmfac9Ph0w",
        "outputId": "be9db7f3-f08a-47f5-d6fa-d7c8bce4f97a"
      },
      "outputs": [],
      "source": [
        "class Provider(StrEnum):\n",
        "    OLLAMA = 'Ollama'\n",
        "\n",
        "clients: dict[Provider, OpenAI] = {}\n",
        "\n",
        "clients[Provider.OLLAMA] = OpenAI(base_url='http://localhost:11434/v1')\n",
        "\n",
        "model = 'llama3.2:latest'\n",
        "client = clients.get(Provider.OLLAMA)\n",
        "if not client:\n",
        "    raise Exception('No client found')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fTHvG2w0sgwU"
      },
      "outputs": [],
      "source": [
        "class Task(StrEnum):\n",
        "  SCAN = 'Scan'\n",
        "  REPORT = 'Threat Report'\n",
        "  PATCH = 'Patch'\n",
        "  TEST = 'Exploit Test'\n",
        "\n",
        "\n",
        "def perform_tasks(tasks, code):\n",
        "  logger.info(f'Performing tasks: {tasks}')\n",
        "\n",
        "  steps = []\n",
        "  if Task.SCAN in tasks:\n",
        "    steps.append('Scan the snippet for security weaknesses and name them clearly.')\n",
        "  if Task.REPORT in tasks:\n",
        "    steps.append('Produce a concise threat report that explains impact, likelihood, and affected components.')\n",
        "  if Task.PATCH in tasks:\n",
        "    steps.append('Propose hardened code that mitigates the identified risks without changing intent.')\n",
        "  if Task.TEST in tasks:\n",
        "    steps.append('Design exploit-style tests or probing steps that would validate the vulnerability.')\n",
        "\n",
        "  task_list = '- ' + '\\n- '.join(steps) if steps else 'No security directive selected.'\n",
        "  system_prompt = f\"\"\"\n",
        "  You are a seasoned application security engineer who recognises languages instantly and\n",
        "  maps weaknesses to common vulnerability classes.\n",
        "  Only rewrite code when instructed to patch it.\n",
        "\n",
        "  Your tasks:\n",
        "  {task_list}\n",
        "  \"\"\"\n",
        "  messages = [\n",
        "      {\"role\": \"system\", \"content\": system_prompt},\n",
        "      {\"role\": \"user\", \"content\": f'Code: \\n{code}'}\n",
        "  ]\n",
        "  response =  client.chat.completions.create(\n",
        "      model=model,\n",
        "      messages=messages\n",
        "  )\n",
        "\n",
        "  content = response.choices[0].message.content\n",
        "\n",
        "  return content"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nlzUyXFus0km"
      },
      "outputs": [],
      "source": [
        "def get_examples() -> tuple[list[any], list[str]]:\n",
        "    python_sql = r'''\n",
        "    import sqlite3\n",
        "\n",
        "    def get_user(conn, user_id):\n",
        "        query = f\"SELECT * FROM users WHERE id = {user_id}\"\n",
        "        cursor = conn.cursor()\n",
        "        cursor.execute(query)\n",
        "        return cursor.fetchone()\n",
        "    '''\n",
        "\n",
        "    js_auth = r'''\n",
        "    app.post('/login', async (req, res) => {\n",
        "        const token = jwt.decode(req.body.token);\n",
        "        if (!token) {\n",
        "            return res.status(401).send('blocked');\n",
        "        }\n",
        "        const user = await db.find(token.user);\n",
        "        res.send(user);\n",
        "    });\n",
        "    '''\n",
        "\n",
        "    go_crypto = r'''\n",
        "    package main\n",
        "\n",
        "    import (\n",
        "        \"crypto/sha1\"\n",
        "    )\n",
        "\n",
        "    func hashPassword(password string) string {\n",
        "        h := sha1.New()\n",
        "        h.Write([]byte(password))\n",
        "        return string(h.Sum(nil))\n",
        "    }\n",
        "    '''\n",
        "\n",
        "    php_upload = r'''\n",
        "    <?php\n",
        "    if ($_FILES[\"file\"][\"error\"] == 0) {\n",
        "        move_uploaded_file($_FILES[\"file\"][\"tmp_name\"], \"/uploads/\" . $_FILES[\"file\"][\"name\"]);\n",
        "        echo \"done\";\n",
        "    }\n",
        "    ?>\n",
        "    '''\n",
        "\n",
        "    rust_config = r'''\n",
        "    use std::env;\n",
        "\n",
        "    fn main() {\n",
        "        let endpoint = env::var(\"SERVICE_URL\").unwrap();\n",
        "        println!(\"Connecting to {}\", endpoint);\n",
        "    }\n",
        "    '''\n",
        "\n",
        "    examples = [\n",
        "        [[Task.SCAN], python_sql, 'python'],\n",
        "        [[Task.REPORT], js_auth, 'javascript'],\n",
        "        [[Task.PATCH], go_crypto, 'go'],\n",
        "        [[Task.TEST], php_upload, 'php'],\n",
        "        [[Task.SCAN, Task.PATCH, Task.REPORT], rust_config, 'rust']\n",
        "    ]\n",
        "\n",
        "    example_labels = [\n",
        "        'Python: SQL injection review',\n",
        "        'JavaScript: Token handling report',\n",
        "        'Go: Strengthen hashing',\n",
        "        'PHP: Exploit upload path',\n",
        "        'Rust: Exposure analysis'\n",
        "    ]\n",
        "\n",
        "    return examples, example_labels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wYReYuvgtDgg"
      },
      "source": [
        "## UI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 664
        },
        "id": "I8Q08SJe8CxK",
        "outputId": "f1d41d06-dfda-4daf-b7ff-6f73bdaf8369"
      },
      "outputs": [],
      "source": [
        "title = 'Security Sentinel'\n",
        "\n",
        "with gr.Blocks(title=title, theme=gr.themes.Monochrome()) as ui:\n",
        "    gr.Markdown(f'# {title}')\n",
        "    gr.Markdown('## Run rapid security sweeps on any snippet.')\n",
        "\n",
        "    with gr.Row():\n",
        "      with gr.Column():\n",
        "        tasks = gr.Dropdown(\n",
        "            label=\"Missions\",\n",
        "            choices=[task.value for task in Task],\n",
        "            value=Task.SCAN,\n",
        "            multiselect=True,\n",
        "            interactive=True,\n",
        "        )\n",
        "        code_input = gr.Code(\n",
        "            label='Code Input',\n",
        "            lines=40,\n",
        "        )\n",
        "        code_language = gr.Textbox(visible=False)\n",
        "\n",
        "      with gr.Column():\n",
        "        gr.Markdown('## Findings')\n",
        "        code_output = gr.Markdown('Awaiting report')\n",
        "\n",
        "\n",
        "    run_btn = gr.Button('Run Audit')\n",
        "\n",
        "    def set_language(tasks, code, language):\n",
        "      syntax_highlights = ['python', 'c', 'cpp', 'javascript', 'typescript', 'go', 'rust', 'php']\n",
        "      logger.debug(f'Tasks: {tasks}, Language: {language}')\n",
        "      highlight = language if language in syntax_highlights else None\n",
        "\n",
        "      return tasks, gr.Code(value=code, language=highlight)\n",
        "\n",
        "    examples, example_labels = get_examples()\n",
        "    examples = gr.Examples(\n",
        "        examples=examples,\n",
        "        example_labels=example_labels,\n",
        "        examples_per_page=20,\n",
        "        inputs=[tasks, code_input, code_language],\n",
        "        outputs=[tasks, code_input],\n",
        "        run_on_click=True,\n",
        "        fn=set_language\n",
        "    )\n",
        "\n",
        "    run_btn.click(perform_tasks, inputs=[tasks, code_input], outputs=[code_output])\n",
        "\n",
        "ui.launch(debug=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "env",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.13.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
